#!/usr/bin/python

##
## Copyright (C) 2011 Luciano Xumerle <luciano.xumerle@gmail.com>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.
##

##
## THE SCRIPT NEEDS TABIX MODULE
##
## http://sourceforge.net/projects/samtools/files/tabix/
##

import sys
import time
import re
import gzip
import tabix

###########################################################################

def getSet( tb, chrom, start ):
   tbresult1 = tb.query(chrom, start,  start+1 )
   set1 = set()
   ensembl = set()
   for a in tbresult1:
       oo=re.sub("^.+gene_name \"", "", a[8])
       oo=re.sub("\".+$", "", oo)
       if oo not in set1:
           set1.add(oo)
       oo=re.sub("^.+gene_id \"", "", a[8])
       oo=re.sub("\".+$", "", oo)
       if oo not in set1:
           ensembl.add(oo)
   return [ set1, ensembl ]

###########################################################################

def readAlignedFile(gtf, fusions, ensembl):
   idx=0
   if ensembl :
       idx=1
   tb = tabix.Tabix(gtf)

   fout = open(fusions, "a")
   for line in open(fusions, "r"):
       poss = line.strip().split("\t")
       chrs=poss[0].split("-")

       set1 = getSet(tb, chrs[0], int(poss[1]) )
       set2 = getSet(tb, chrs[1], int(poss[2]) )

       print " + ".join(list(set1[idx])) + "\t" + " + ".join(list(set2[idx])) + "\t" + line.strip()

###########################################################################

def main():
   """main routine"""

   # pars
   database=''
   fusions=''
   ensembl = False
   errorMsg="SYNTAX: " + sys.argv[0] + " [-ensCode] <tabixGTF> <fusions.out>\n"

   if len(sys.argv)<2:
       print >>sys.stderr, errorMsg
       return
   for a in sys.argv[1:]:
       if a == "-ensCode" :
           ensembl = True
       if database=='':
           database=a
           continue
       if fusions=='':
           fusions=a
           continue

   if database=='' or fusions=='':
       print >>sys.stderr, errorMsg
       return

   print >>sys.stderr, time.ctime() + " START COMPUTATION..."

   readAlignedFile(database, fusions, ensembl)

   print >>sys.stderr, time.ctime() + " DONE"

if __name__=='__main__': main()
